#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

template<typename TYPE_START, typename TYPE_END, typename TYPE_WEIGHT, typename TYPE_Y>
class KernelLerp
{
public:
    __aicore__ inline KernelLerp() {}
    __aicore__ inline void Init(GM_ADDR start, GM_ADDR end, GM_ADDR weight, GM_ADDR y, uint32_t ALIGN_NUM, 
                                uint32_t totalLength, uint32_t core_size, uint32_t core_remain, uint32_t block_size, uint32_t weight_length)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!")
        this->blockLength = core_size + (GetBlockNum() == GetBlockIdx() + 1 ? core_remain : 0);
        this->tileLength = block_size;
        this->blockLength = this->blockLength + (this->blockLength % ALIGN_NUM ? ALIGN_NUM - this->blockLength % ALIGN_NUM : 0);

        this->wSize = weight_length;

        auto startPointer = core_size * GetBlockIdx();
        auto bufferlength = this->blockLength;


        sGm.SetGlobalBuffer((__gm__ TYPE_START *)start + startPointer, bufferlength);
        eGm.SetGlobalBuffer((__gm__ TYPE_END *)end + startPointer, bufferlength);
        yGm.SetGlobalBuffer((__gm__ TYPE_Y *)y + startPointer, bufferlength);
        wGm.SetGlobalBuffer((__gm__ TYPE_WEIGHT *)weight + startPointer, bufferlength);

        this->tileNum = this->blockLength / this->tileLength + (this->blockLength % this->tileLength > 0);

        pipe.InitBuffer(sQueue, BUFFER_NUM, this->tileLength * sizeof(TYPE_START));
        pipe.InitBuffer(eQueue, BUFFER_NUM, this->tileLength * sizeof(TYPE_END));
        pipe.InitBuffer(wQueue, BUFFER_NUM, this->tileLength * sizeof(TYPE_WEIGHT));
        pipe.InitBuffer(yQueue, BUFFER_NUM, this->tileLength * sizeof(TYPE_Y));
        pipe.InitBuffer(tmp1, this->tileLength * sizeof(float));
        pipe.InitBuffer(tmp2, this->tileLength * sizeof(float));
        // this->weight = wGm.GetValue(0);
    }


    __aicore__ inline void Process() 
    {
        int32_t loopCount = this->tileNum;
        for (int32_t i = 0; i < loopCount-1; i++) 
        {
            CopyIn(i, this->tileLength);
            Compute(i, this->tileLength);
            CopyOut(i, this->tileLength);
        }

        uint32_t length = this->blockLength - this->tileLength * (loopCount - 1);
        CopyIn(loopCount - 1, length);
        Compute(loopCount - 1, length);
        CopyOut(loopCount - 1, length);
    }


private:
    __aicore__ inline void CopyIn(int32_t progress, uint32_t length) {
        LocalTensor<TYPE_START> sLocal = sQueue.AllocTensor<TYPE_START>();
        LocalTensor<TYPE_END> eLocal = eQueue.AllocTensor<TYPE_END>();
        LocalTensor<TYPE_WEIGHT> wLocal = wQueue.AllocTensor<TYPE_WEIGHT>();

        DataCopy(sLocal, sGm[progress * this->tileLength], length);
        DataCopy(eLocal, eGm[progress * this->tileLength], length);
        DataCopy(wLocal, wGm[progress * this->tileLength], length);

        sQueue.EnQue(sLocal);
        eQueue.EnQue(eLocal);
        wQueue.EnQue(wLocal);

    }

    __aicore__ inline void Compute(int32_t progress, uint32_t length)
    {
        LocalTensor<TYPE_START> sLocal = sQueue.DeQue<TYPE_START>();
        LocalTensor<TYPE_END> eLocal = eQueue.DeQue<TYPE_END>();
        LocalTensor<TYPE_WEIGHT> wLocal = wQueue.DeQue<TYPE_WEIGHT>();

        LocalTensor<TYPE_Y> yLocal = yQueue.AllocTensor<TYPE_Y>();


        if constexpr (std::is_same_v<TYPE_Y, half>)
        {
            LocalTensor<float> p1 = tmp1.Get<float>();
            LocalTensor<float> p2 = tmp2.Get<float>();
            
            Cast(p1, sLocal, RoundMode::CAST_NONE, length);
            Cast(p2, eLocal, RoundMode::CAST_NONE, length);
            Sub(p2, p2, p1, length);
            if (wSize == 1)
            {
                Muls(p2, p2, (float)wGm.GetValue(0), length);
            }
            else
            {
                Cast(p1, wLocal, RoundMode::CAST_NONE, length);
                Mul(p2, p1, p2, length);
            }
            Cast(p1, sLocal, RoundMode::CAST_NONE, length);
            Add(p2, p2, p1, length);
            Cast(yLocal, p2, RoundMode::CAST_NONE, length);
        }
        else
        {
            Sub(eLocal, eLocal, sLocal, length);
            if (wSize == 1)
            {
                Muls(eLocal, eLocal, wGm.GetValue(0), length);
            }
            else
            {
                Mul(eLocal, eLocal, wLocal, length);
            }     
            Add(yLocal, sLocal, eLocal, length);
        }

        sQueue.FreeTensor(sLocal);
        eQueue.FreeTensor(eLocal);
        wQueue.FreeTensor(wLocal);
        yQueue.EnQue<TYPE_Y>(yLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress, uint32_t length) {
        LocalTensor<TYPE_Y> yLocal = yQueue.DeQue<TYPE_Y>();
        DataCopy(yGm[progress * this->tileLength], yLocal, length);
        yQueue.FreeTensor(yLocal);
    }

    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> sQueue, eQueue, wQueue;
    TQue<QuePosition::VECOUT, BUFFER_NUM> yQueue;
    GlobalTensor<TYPE_START> sGm;
    GlobalTensor<TYPE_END> eGm;
    GlobalTensor<TYPE_WEIGHT> wGm;
    GlobalTensor<TYPE_Y> yGm;
    TBuf<QuePosition::VECCALC> tmp1, tmp2;

    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t wSize;
};


template<typename TYPE_START, typename TYPE_END, typename TYPE_WEIGHT, typename TYPE_Y>
class KernelLerp_Broadcast
{
public:
    __aicore__ inline KernelLerp_Broadcast() {}
    __aicore__ inline void Init(GM_ADDR start, GM_ADDR end, GM_ADDR weight, GM_ADDR y, uint32_t ALIGN_NUM, 
            uint32_t totalLength, uint32_t start_length, uint32_t end_length, uint32_t weigth_length, 
            uint32_t core_size, uint32_t core_remain, uint32_t block_size)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!")

        this->startLength = start_length;
        this->endLength = end_length;
        this->weightLength = weigth_length;

        this->blockLength = core_size + (GetBlockNum() == GetBlockIdx() + 1 ? core_remain : 0);
        this->tileLength = block_size;
        this->blockLength = this->blockLength + (this->blockLength % ALIGN_NUM ? ALIGN_NUM - this->blockLength % ALIGN_NUM : 0);

        this->startPointer = core_size * GetBlockIdx();
        auto bufferlength = this->blockLength;


        sGm.SetGlobalBuffer((__gm__ TYPE_START *)start, bufferlength);
        eGm.SetGlobalBuffer((__gm__ TYPE_END *)end, bufferlength);
        yGm.SetGlobalBuffer((__gm__ TYPE_Y *)y, bufferlength);
        wGm.SetGlobalBuffer((__gm__ TYPE_WEIGHT *)weight, bufferlength);

        this->tileNum = this->blockLength / this->tileLength + (this->blockLength % this->tileLength > 0);

        pipe.InitBuffer(sQueue, BUFFER_NUM, this->tileLength * sizeof(TYPE_START));
        pipe.InitBuffer(eQueue, BUFFER_NUM, this->tileLength * sizeof(TYPE_END));
        pipe.InitBuffer(wQueue, BUFFER_NUM, this->tileLength * sizeof(TYPE_WEIGHT));
        pipe.InitBuffer(yQueue, BUFFER_NUM, this->tileLength * sizeof(TYPE_Y));
        pipe.InitBuffer(tmp1, this->tileLength * sizeof(float));
        pipe.InitBuffer(tmp2, this->tileLength * sizeof(float));
    }

    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum;
        for (int32_t i = 0; i < loopCount-1; i++) 
        {
            uint32_t position = startPointer + i * this->tileLength;
            CopyIn(position, this->tileLength);
            Compute(position, this->tileLength);
            CopyOut(position, this->tileLength);
        }
        uint32_t position = startPointer + (loopCount - 1) * this->tileLength;
        uint32_t length = this->blockLength - this->tileLength * (loopCount - 1);
        CopyIn(position, length);
        Compute(position, length);
        CopyOut(position, length);
    }


private:
    __aicore__ inline void CopyIn(int32_t position, uint32_t length) {
        LocalTensor<TYPE_START> sLocal = sQueue.AllocTensor<TYPE_START>();
        LocalTensor<TYPE_END> eLocal = eQueue.AllocTensor<TYPE_END>();
        LocalTensor<TYPE_WEIGHT> wLocal = wQueue.AllocTensor<TYPE_WEIGHT>();

        DataCopy(sLocal, sGm[position % startLength], length);
        DataCopy(eLocal, eGm[position % endLength], length);
        DataCopy(wLocal, wGm[position % weightLength], length);

        sQueue.EnQue(sLocal);
        eQueue.EnQue(eLocal);
        wQueue.EnQue(wLocal);

    }

    __aicore__ inline void Compute(int32_t position, uint32_t length)
    {
        LocalTensor<TYPE_START> sLocal = sQueue.DeQue<TYPE_START>();
        LocalTensor<TYPE_END> eLocal = eQueue.DeQue<TYPE_END>();
        LocalTensor<TYPE_WEIGHT> wLocal = wQueue.DeQue<TYPE_WEIGHT>();

        LocalTensor<TYPE_Y> yLocal = yQueue.AllocTensor<TYPE_Y>();


        if constexpr (std::is_same_v<TYPE_Y, half>)
        {
            LocalTensor<float> p1 = tmp1.Get<float>();
            LocalTensor<float> p2 = tmp2.Get<float>();
            
            Cast(p1, sLocal, RoundMode::CAST_NONE, length);
            Cast(p2, eLocal, RoundMode::CAST_NONE, length);
            Sub(p2, p2, p1, length);
            if (weightLength == 1)
            {
                Muls(p2, p2, (float)wGm.GetValue(0), length);
            }
            else
            {
                Cast(p1, wLocal, RoundMode::CAST_NONE, length);
                Mul(p2, p1, p2, length);
            }
            Cast(p1, sLocal, RoundMode::CAST_NONE, length);
            Add(p2, p2, p1, length);
            Cast(yLocal, p2, RoundMode::CAST_NONE, length);
        }
        else
        {
            Sub(eLocal, eLocal, sLocal, length);
            if (weightLength == 1)
            {
                Muls(eLocal, eLocal, wGm.GetValue(0), length);
            }
            else
            {
                Mul(eLocal, eLocal, wLocal, length);
            }
            Add(yLocal, sLocal, eLocal, length);
        }


        // Adds(yLocal, wLocal, (TYPE_END)0.0, length);
        sQueue.FreeTensor(sLocal);
        eQueue.FreeTensor(eLocal);
        wQueue.FreeTensor(wLocal);
        yQueue.EnQue<TYPE_Y>(yLocal);
    }

    __aicore__ inline void CopyOut(int32_t position, uint32_t length) {
        LocalTensor<TYPE_Y> yLocal = yQueue.DeQue<TYPE_Y>();
        DataCopy(yGm[position], yLocal, length);
        yQueue.FreeTensor(yLocal);
    }

    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> sQueue, eQueue, wQueue;
    TQue<QuePosition::VECOUT, BUFFER_NUM> yQueue;
    GlobalTensor<TYPE_START> sGm;
    GlobalTensor<TYPE_END> eGm;
    GlobalTensor<TYPE_WEIGHT> wGm;
    GlobalTensor<TYPE_Y> yGm;
    TBuf<QuePosition::VECCALC> tmp1, tmp2;

    // TYPE_WEIGHT weight;
    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t startPointer;
    uint32_t startLength;
    uint32_t endLength;
    uint32_t weightLength;
};



extern "C" __global__ __aicore__ void lerp(GM_ADDR start, GM_ADDR end, GM_ADDR weight, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    
    if (TILING_KEY_IS(1))
    {
        KernelLerp<DTYPE_START, DTYPE_END, DTYPE_WEIGHT, DTYPE_Y> op;
        op.Init(start, end, weight, y,
            tiling_data.ALIGN_NUM, tiling_data.total_length, tiling_data.core_size, tiling_data.core_remain, tiling_data.block_size, tiling_data.weight_length);
        op.Process();
    }
    else if(TILING_KEY_IS(2))
    {
        KernelLerp_Broadcast<DTYPE_START, DTYPE_END, DTYPE_WEIGHT, DTYPE_Y> op;
        op.Init(start, end, weight, y,
            tiling_data.ALIGN_NUM, tiling_data.total_length, tiling_data.start_length, tiling_data.end_length, tiling_data.weight_length, tiling_data.core_size, tiling_data.core_remain, tiling_data.block_size);
        op.Process();
    }
    
}
